import cv2
import os
import numpy as np


class FaceCascader:
    def __init__(self):
        # 加载正脸检测分类器
        self.frontal_face_cascade = cv2.CascadeClassifier(
            cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

        # 加载侧脸检测分类器
        self.profile_face_cascade = cv2.CascadeClassifier(
            cv2.data.haarcascades + 'haarcascade_profileface.xml')

    def to_gray(self, img: np.ndarray):
        return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    def detect_frontal_faces(self, img: np.ndarray):
        # 检测正脸
        return self.frontal_face_cascade.detectMultiScale(img, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

    def detect_profile_faces(self, img: np.ndarray):
        # 检测侧脸
        return self.profile_face_cascade.detectMultiScale(img, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))


def detect(img_path: str, res_dir: str):
    img_basename = os.path.basename(img_path)

    faceCascader = FaceCascader()
    image = cv2.imread(img_path)
    gray_img = faceCascader.to_gray(image)

    frontal_faces = faceCascader.detect_frontal_faces(gray_img)
    profile_faces = faceCascader.detect_profile_faces(gray_img)

    frontal_res = image.copy()
    for (x, y, w, h) in frontal_faces:
        cv2.rectangle(frontal_res, (x, y), (x + w, y + h), (255, 0, 0), 2)
    cv2.imwrite(f"{res_dir}/frontal_faces_{img_basename}", frontal_res)

    profile_res = image
    for (x, y, w, h) in profile_faces:
        cv2.rectangle(profile_res, (x, y), (x + w, y + h), (0, 255, 0), 2)
    cv2.imwrite(f"{res_dir}/profile_faces_{img_basename}", profile_res)


def main():
    for img_path in os.listdir("assets"):
        if img_path.lower().endswith("jpg"):
            detect(f"assets/{img_path}", "result")


if __name__ == "__main__":
    main()
